#pragma once

#include <string>
#include <map>
#include <utility>
#include <optional>

// #include "txdnn_internal.h"
#include "macros/errors.h"
#include "operators/invoker.h"


namespace txdnn {

/* class InvokerCache */
class InvokerCache {
public:
  // network_config -> id
  using Key = std::pair<std::string, std::string>;

  std::optional<txdnn::Invoker> operator[](const Key& key) const {
    const auto cfg_item = cfg_item_map_.find(key.first);
    if (cfg_item == cfg_item_map_.end()) {
      return std::nullopt;
    }

    const auto& id_invoker  = cfg_item->second.id_invoker_map_;
    const auto invoker      = id_invoker.find(key.second);
    if (invoker == id_invoker.end()) {
      return std::nullopt;
    }
    return invoker->second;
  }

  std::optional<Invoker> getNameId(const std::string& cfg,
                                   const std::string& name) const {
    const auto cfg_item = cfg_item_map_.find(cfg);

    if (cfg_item == cfg_item_map_.end()) {
      // TXLOG_DEBUG("No invoker found for %s \n", cfg.c_str());
      return std::nullopt;
    }
    if (cfg_item->second.name_id_map_.empty()) {
      // TXLOG_DEBUG("Invoker found for %s, but no name_id_map_.\n", cfg.c_str());
      return std::nullopt;
    }

    const auto& id_invoker  = cfg_item->second.id_invoker_map_;
    const auto& name_ids    = cfg_item->second.name_id_map_;
    const auto name_id      = name_ids.find(name);

    if (name_id == name_ids.end()) {
      // TXLOG_DEBUG("Invoker found for %s, but no one match an algorithm %s \n", cfg.c_str(), name.c_str());
      return std::nullopt;
    }

    const auto id_invoker_it = id_invoker.find(name_id->second);
    if (id_invoker_it == id_invoker.end()) {
      TXDNN_THROW("No invoker found, no one match the network config \n");
    }
    return id_invoker_it->second;
  }

  std::optional<std::string> getNameIdOpId(const std::string& cfg,
                                           const std::string& algorithm) const {
    const auto item_it = cfg_item_map_.find(cfg);

    if (item_it == cfg_item_map_.end()) {
      // TXLOG_DEBUG("No invoker found for %s", cfg.c_str());
      return std::nullopt;
    }

    if (item_it->second.name_id_map_.empty()) {
      // TXLOG_DEBUG("Invoker found for %s, but no name_id_map_.", cfg.c_str());
      return std::nullopt;
    }

    const auto& name_ids  = item_it->second.name_id_map_;
    const auto name_id    = name_ids.find(algorithm);

    if (name_id == name_ids.end()) {
      // TXLOG_DEBUG("Invoker found for %s, but no one match an algorithm %s \n", cfg.c_str(), algorithm.c_str());
      return std::nullopt;
    }

    return name_id->second;
  }

  void registerInvoker(const Key& key, const Invoker& invoker) {
    auto it = cfg_item_map_.find(key.first);

    if (it != cfg_item_map_.end()) {
      it->second.id_invoker_map_[key.second] = invoker;
    } else {
      auto& item = cfg_item_map_.insert({key.first, Item{}}).first->second;
      item.id_invoker_map_[key.second] = invoker;
    }
    // TXLOG_DEBUG("Register invoker for algorithm %s and operator %s\n", key.first.c_str(), key.second.c_str());
  }

  void setNameAndId(const std::string& cfg,
                   const std::string& name,
                   const std::string& id) {
    const auto item_it = cfg_item_map_.find(cfg);
    if (item_it == cfg_item_map_.end()) {
      TXDNN_THROW("No invoker was registered for " + cfg);
    }

    {
      const auto& id_invoker_map = item_it->second.id_invoker_map_;
      const auto invoker_it      = id_invoker_map.find(id);
      if (invoker_it == id_invoker_map.end()) {
        TXDNN_THROW("No invoker with id of " + id + " was registered for " + cfg);
      }
    }

    item_it->second.name_id_map_[name] = id;
    // TXLOG_DEBUG("Register invoker for id, algorithm %s and operator %s \n", id.c_str(), name.c_str(), cfg.c_str());
  }

private:
  struct Item {
    // algorithm_name -> id
    std::map<std::string, std::string> name_id_map_;
    // id -> invoker
    std::map<std::string, Invoker> id_invoker_map_;
  };
  // network_config -> item
  std::map<std::string, Item> cfg_item_map_;
};

} // namespace txdnn